import hashlib
import json

import pandas as pd
from pandas.core.interchange.dataframe_protocol import DataFrame

tag_to_type = {
        1: "疾病和诊断",
        2: "疾病和诊断",
        3: "检查",
        4: "检查",
        5: "检验",
        6: "检验",
        7: "手术",
        8: "手术",
        9: "药物",
        10: "药物",
        11: "解剖部位",
        12: "解剖部位",
        # 可根据实际标签扩展更多类型
    }

schema = {
    "疾病和诊断": None,
    "检查": None,
    "检验": None,
    "手术": None,
    "药物": None,
    "解剖部位": None
}

def gen_md5(text):
    # 创建 MD5 对象
    md5 = hashlib.md5()

    # 更新 MD5 对象（需先将字符串编码为 bytes）
    md5.update(text.encode('utf-8'))

    # 获取十六进制表示的 MD5 值
    md5_hex = md5.hexdigest()
    return md5_hex

def convert_to_target_format(tokens, ner_tags, example_id="c742d5a49dc5db8b5b1f577f9e16e2f4"):
    # ----------------------
    # 1. 拼接文本（text字段）
    # ----------------------
    text = ''.join(tokens)  # 合并连续空格（可选处理）

    # ----------------------
    # 2. 提取实体（info_list字段）
    # ----------------------
    info_list = []

    # 计算每个token的起始和结束偏移量（基于字符数，全角符号占2字符，半角占1字符）
    current_entity = None  # 用于跟踪当前正在处理的实体
    # 遍历tokens和标签，提取实体
    for idx, (token, tag) in enumerate(zip(tokens, ner_tags)):
        # print(f"idx: {idx}, token: {token}, tag: {tag}")
        if tag == 0:
            # 标签为0时，结束当前实体（如果存在）
            if current_entity:
                info_list.append([current_entity])
                current_entity = None
            continue

        # 获取实体类型（若标签不在映射中，则跳过）
        entity_type = tag_to_type.get(tag, None)
        if not entity_type:
            if current_entity:
                info_list.append([current_entity])
                current_entity = None
            continue

        # 初始化或扩展当前实体
        if current_entity:
            # 合并连续同类型实体
            if current_entity["type"] == entity_type:
                current_entity["span"] += token
                # 更新偏移量（当前token的结束偏移为下一个token的开始偏移）
                current_entity["offset"][1] += 1
            else:
                # 类型变更，结束上一个实体
                info_list.append([current_entity])
                current_entity = {
                    "type": entity_type,
                    "span": token,
                    "offset": [idx, idx + 1]
                }
        else:
            current_entity = {
                "type": entity_type,
                "span": token,
                "offset": [idx, idx + 1]
            }

    # 处理最后一个实体
    if current_entity:
        info_list.append([current_entity])

    return text, info_list


# ----------------------
# 示例调用
# ----------------------
if __name__ == "__main__":
    # 方法二（推荐大文件）：指定引擎为 pyarrow
    # parquet_file = 'train-00000-of-00001.parquet'
    # target_file = 'train.csv'
    parquet_file = 'validation-00000-of-00001.parquet'
    target_file = 'dev.csv'
    df = pd.read_parquet(parquet_file, engine='pyarrow')

    data = {
        'id': [],
        'text': [],
        'info_list': [],
        'schema': []
    }

    for row in df.itertuples(index=True, name='Person'):
        text, info_list = convert_to_target_format(row.tokens, row.ner_tags)
        print(f"text: {text}, info list: {info_list}")
        data['id'].append(gen_md5(text))
        data['text'].append(text)
        data['info_list'].append(json.dumps(info_list, ensure_ascii=False))
        data['schema'].append(json.dumps(schema, ensure_ascii=False))

    data_df = pd.DataFrame(data)
    data_df.to_csv(target_file, index=False)
